module aru_reduce_stage2 (
    input logic                 clk,
    input logic                 rst_n,
          aru_reduce_ctrl_if.in u_aru_ctrl_if,
          aru_reduce_pld_if.in  u_aru_reduce_stage1_if,
          aru_reduce_pld_if.out u_aru_reduce_stage2_if
);

    // ===== 暂存控制信号 =====
    logic       reduce_m_reg;
    logic       reduce_n_reg;
    logic [1:0] reduce_op_reg;

    // 配置信号握手逻辑
    logic ctrl_rdy, ctrl_vld;
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            ctrl_rdy      <= 1'b1;
            reduce_m_reg  <= 'd0;
            reduce_n_reg  <= 'd0;
            reduce_op_reg <= 'd0;
        end else if (ctrl_rdy == 1'b0) begin
            if (u_aru_reduce_stage2_if.vld && u_aru_reduce_stage2_if.rdy && u_aru_reduce_stage2_if.sdb.eom && u_aru_reduce_stage2_if.sdb.eon) begin
                ctrl_rdy <= 1'b1;
            end
        end else begin
            if (u_aru_ctrl_if.vld && u_aru_ctrl_if.rdy) begin
                ctrl_rdy      <= 1'b0;
                reduce_m_reg  <= u_aru_ctrl_if.reduce_m;
                reduce_n_reg  <= u_aru_ctrl_if.reduce_n;
                reduce_op_reg <= u_aru_ctrl_if.reduce_op;
            end
        end
    end

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            ctrl_vld <= 1'b0;
        end else if (ctrl_vld == 1'b0) begin
            if (u_aru_ctrl_if.vld && u_aru_ctrl_if.rdy) begin
                ctrl_vld <= 1'b1;
            end
        end else begin
            if (u_aru_reduce_stage1_if.vld && u_aru_reduce_stage1_if.rdy && u_aru_reduce_stage1_if.sdb.eom && u_aru_reduce_stage1_if.sdb.eon) begin
                ctrl_vld <= 1'b0;
            end
        end
    end

    logic end_reduce;  // 指示一次累*任务的输出/结束
    assign end_reduce = u_aru_reduce_stage1_if.sdb.eom && u_aru_reduce_stage1_if.sdb.eon;

    // 流水线使能信号
    logic psum_en;
    assign psum_en = ~reduce_m_reg && reduce_n_reg && u_aru_reduce_stage1_if.vld;
    logic pipeline_en;
    assign pipeline_en = u_aru_reduce_stage2_if.rdy;
    logic bypass_en;
    always_comb begin
        case ({
            reduce_m_reg, reduce_n_reg
        })
            2'b00: bypass_en = 1'b1;
            2'b01: bypass_en = 1'b0;
            2'b10: bypass_en = 1'b1;
            2'b11: bypass_en = 1'b0;
        endcase
    end

    // ==================== 4个 Max/Min 模块例化 ====================
    // 每组处理 2个输入（8个输入减少到4个输出）
    localparam int NUM_GROUPS = 4;
    localparam int INPUTS_PER_GROUP = 2;

    bf16_t maxmin_dat_out    [NUM_GROUPS];
    logic  maxmin_dat_out_vld[NUM_GROUPS];

    genvar i;
    generate
        for (i = 0; i < NUM_GROUPS; i++) begin : gen_maxmin
            logic  dat_in_vld[1:0];
            bf16_t dat_in    [1:0];

            always_comb begin
                for (int j = 0; j < INPUTS_PER_GROUP; j++) begin
                    dat_in_vld[j] = u_aru_reduce_stage1_if.msk[i*INPUTS_PER_GROUP+j];
                    dat_in[j]     = u_aru_reduce_stage1_if.dat[i*INPUTS_PER_GROUP+j].bf16.dat;
                end
            end

            aru_reduce_max_min_2_to_1 u_maxmin (
                .clk        (clk),
                .rst_n      (rst_n),
                .en         (pipeline_en),
                .is_max     (reduce_op_reg == 2'b00),
                .psum_en    (psum_en),
                .end_reduce (end_reduce),
                .dat_in_vld (dat_in_vld),
                .dat_in     (dat_in),
                .dat_out    (maxmin_dat_out[i]),
                .dat_out_vld(maxmin_dat_out_vld[i])
            );
        end
    endgenerate

    // ==================== 4个 Sum 模块例化 ====================
    fp32_t sum_dat_out    [NUM_GROUPS];
    logic  sum_dat_out_vld[NUM_GROUPS];

    generate
        for (i = 0; i < NUM_GROUPS; i++) begin : gen_sum
            logic  dat_in_vld[1:0];
            fp32_t dat_in    [1:0];
            always_comb begin
                for (int j = 0; j < INPUTS_PER_GROUP; j++) begin
                    dat_in_vld[j] = u_aru_reduce_stage1_if.msk[i*INPUTS_PER_GROUP+j];
                    dat_in[j]     = u_aru_reduce_stage1_if.dat[i*INPUTS_PER_GROUP+j].fp32;
                end
            end

            aru_reduce_sum_2_to_1 u_sum (
                .clk        (clk),
                .rst_n      (rst_n),
                .en         (pipeline_en),
                .psum_en    (psum_en),
                .end_reduce (end_reduce),
                .dat_in_vld (dat_in_vld),
                .dat_in     (dat_in),
                .dat_out    (sum_dat_out[i]),
                .dat_out_vld(sum_dat_out_vld[i])
            );
        end
    endgenerate

    // ==================== 输出选择和流水线 ====================
    // 根据 reduce_op 选择输出
    // reduce_op: 0: max, 1: min, 2: sum, 3: avg (avg与sum相同处理)

    bf16_fp32_u output_dat[NUM_GROUPS];
    logic       output_vld[NUM_GROUPS];

    always_comb begin
        for (int i = 0; i < NUM_GROUPS; i++) begin
            if (reduce_op_reg == 2'b00 || reduce_op_reg == 2'b01) begin
                // Max 或 Min 操作
                output_dat[i].bf16.dat = maxmin_dat_out[i];
                output_dat[i].bf16.rsv = '0;
                output_vld[i]          = maxmin_dat_out_vld[i];
            end else begin
                // Sum 或 Avg 操作
                output_dat[i].fp32 = sum_dat_out[i];
                output_vld[i]      = sum_dat_out_vld[i];
            end
        end
    end

    // ==================== Bypass模式（reduce_m=0 & reduce_n=0）====================

    // 6级延迟寄存器（两种模式共用，用于sideband打拍）
    logic            pipe_vld[5:0];  // 单独的 valid 信号数组
    aru_reduce_pld_t pipe_pld[5:0];  // 数据结构体数组

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            for (int i = 0; i < 6; i++) begin
                pipe_vld[i] <= 1'b0;
                pipe_pld[i] <= '0;
            end
        end else if (pipeline_en) begin
            // Stage 0: 输入采样
            pipe_vld[0]     <= u_aru_reduce_stage1_if.vld;
            pipe_pld[0].dat <= u_aru_reduce_stage1_if.dat;
            pipe_pld[0].msk <= u_aru_reduce_stage1_if.msk;
            pipe_pld[0].sdb <= u_aru_reduce_stage1_if.sdb;

            for (int i = 1; i < 6; i++) begin
                pipe_vld[i] <= pipe_vld[i-1];
                pipe_pld[i] <= pipe_pld[i-1];
            end
        end
    end

    // 输出寄存器 - 添加一级流水，确保时序
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            u_aru_reduce_stage2_if.vld <= 1'b0;
            u_aru_reduce_stage2_if.dat <= '0;
            u_aru_reduce_stage2_if.msk <= '0;
            u_aru_reduce_stage2_if.sdb <= '0;
        end else if (pipeline_en) begin
            // sideband和vld统一使用打拍后的数据
            u_aru_reduce_stage2_if.vld <= pipe_vld[4];
            u_aru_reduce_stage2_if.sdb <= pipe_pld[4].sdb;

            if (bypass_en) begin
                // Bypass模式：输出与输入一致，延迟4拍
                u_aru_reduce_stage2_if.dat <= pipe_pld[4].dat;
                u_aru_reduce_stage2_if.msk <= pipe_pld[4].msk;
            end else begin
                // Reduce模式：输出4个结果（从8个输入减少到4个输出）
                // 输出4个结果到低4位
                for (int i = 0; i < NUM_GROUPS; i++) begin
                    u_aru_reduce_stage2_if.dat[i] <= output_dat[i];
                    u_aru_reduce_stage2_if.msk[i] <= output_vld[i];
                end
                // 高28位清零
                for (int i = NUM_GROUPS; i < `P_ARU * `N0; i++) begin
                    u_aru_reduce_stage2_if.dat[i] <= '0;
                    u_aru_reduce_stage2_if.msk[i] <= '0;
                end
            end
        end
    end

    // 配置接口握手
    assign u_aru_ctrl_if.rdy          = ctrl_rdy;

    // Stage1 接口握手
    assign u_aru_reduce_stage1_if.rdy = u_aru_reduce_stage2_if.rdy & ctrl_vld;

endmodule
